import pandas as pd
import numpy as np
import torch
import networkx as nx
from itertools import combinations
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.compose import ColumnTransformer


def get_data(csv_path="timestamp_label2.csv"):
    """
    读取CSV，返回:
        node_features : torch.FloatTensor [N, F]
        adj_matrix    : torch.FloatTensor [N, N]
        targets       : np.ndarray        [N]
    """
    #  读 CSV
    df = pd.read_csv(csv_path, parse_dates=["time"])

    # 建 node 映射
    node_ids = df["topicID"].unique()
    id2idx = {tid: idx for idx, tid in enumerate(node_ids)}

    # 生成节点特征
    numeric_cols = ["hot_value", "rank_top"]  # 数值
    categorical_cols = ["platform", "cls"]  # 类别
    # 编码 + 标准化
    pre = ColumnTransformer([
        ("num", StandardScaler(), numeric_cols),
        ("cat", OneHotEncoder(handle_unknown="ignore"), categorical_cols)
    ])
    X = pre.fit_transform(df[numeric_cols + categorical_cols])
    X = X.toarray() if hasattr(X, "toarray") else X  # 稀疏转稠密

    # 聚合到 topicID 维度（取平均）
    feat = pd.DataFrame(X)
    feat["topicID"] = df["topicID"].values
    F = feat.shape[1] - 1
    node_features = torch.zeros((len(node_ids), F), dtype=torch.float32)
    for tid, g in feat.groupby("topicID"):
        node_features[id2idx[tid]] = torch.tensor(
            g.drop(columns=["topicID"]).mean().values,
            dtype=torch.float32
        )

    # 构建共现图
    g = nx.Graph()
    g.add_nodes_from(range(len(node_ids)))
    for _, grp in df.groupby("timestamp"):
        idxs = [id2idx[tid] for tid in grp["topicID"]]
        for u, v in combinations(idxs, 2):
            # 多次共现就累加权重
            if g.has_edge(u, v):
                g[u][v]["weight"] += 1
            else:
                g.add_edge(u, v, weight=1)

    num_nodes = len(node_ids)
    adj_matrix = torch.zeros((num_nodes, num_nodes), dtype=torch.float32)
    for u, v, d in g.edges(data=True):
        adj_matrix[u, v] = adj_matrix[v, u] = d["weight"]

    # 目标值
    targets = df.groupby("topicID")["hot_value"].mean().values

    return node_features, adj_matrix, targets
